"""
    create by IAmFiveHigh on 2023/8/25
"""
from .Matrix import Matrix
from .Vector import Vector
from .__globals import is_zero


class LinearSystem:
    def __init__(self, A: Matrix, b: Matrix or Vector = None):
        self._m = A.row_num()
        self._n = A.col_num()

        if b is not None and self._m != len(b):
            raise Exception(f"Error: {A}行数[{self._m}]应该与{b}列数[{len(b)}]相等")

        if b is None:
            self.Ab = [Vector(A.row_vector(i).under_list()) for i in range(self._m)]

        if isinstance(b, Matrix):
            self.Ab = [Vector(A.row_vector(i).under_list() + b.row_vector(i).under_list())
                       for i in range(self._m)]
        elif isinstance(b, Vector):
            self.Ab = [Vector(A.row_vector(i).under_list() + [b[i]])
                       for i in range(self._m)]

        #  保存主元的列
        self.pivots = []

    def _max_row(self, index_i, index_j, n):
        # 寻找主元列最大值
        best, ret = abs(self.Ab[index_i][index_j]), index_i
        for i in range(index_i + 1, n):
            if abs(self.Ab[i][index_j]) > best:
                best, ret = abs(self.Ab[i][index_j]), i
        return ret

    def _forward(self):
        # 前向消元
        i, k = 0, 0
        while i < self._m and k < self._n:
            # 获取K列最大值的行
            max_row = self._max_row(i, k, self._m)
            # 交换主元列最大值的行和当前行
            self.Ab[i], self.Ab[max_row] = self.Ab[max_row], self.Ab[i]

            if is_zero(self.Ab[i][k]):
                # 如果k列最大值都为0，说明整个k列都为0，寻找下一列
                k += 1
            else:
                # 将主元行除以自己主元的值使主元值 == 1
                self.Ab[i] = self.Ab[i] / self.Ab[i][k]
                for j in range(i + 1, self._m):
                    # 其他行减去主元行的倍数使其他行主元列值 == 0
                    self.Ab[j] = self.Ab[j] - self.Ab[j][k] * self.Ab[i]
                # 保存主元所在列数
                self.pivots.append(k)
                i += 1

    def _backward(self):
        # 后向消元
        n = len(self.pivots)
        for i in range(n - 1, -1, -1):
            k = self.pivots[i]
            # Ab[i][k]为主元
            for j in range(i - 1, -1, -1):
                self.Ab[j] = self.Ab[j] - self.Ab[j][k] * self.Ab[i]

    def gauss_jordan_elimination(self) -> bool:
        """如果有解，返回True；如果没有解，返回False"""
        self._forward()
        self._backward()

        for i in range(len(self.pivots), self._m):
            # 从最后一个主元行往下查看
            if not is_zero(self.Ab[i][-1]):
                # 如果最右边的值不等于0 形成0=b局面, 方程组无解
                return False
        return True

    def fancy_print(self):

        for i in range(self._m):
            # temp_array把原向量中为0但是显示为极小数的数字变成0
            temp_array = list(map(LinearSystem._fancy_number, self.Ab[i].under_list()))

            right = len(self.Ab[0]) - self._n
            print(" ".join(str(temp_array[j]) for j in range(self._n)), end=" ")
            print("|", end=" ")
            print(" ".join(str(temp_array[j]) for j in range(-right, 0)))

    @staticmethod
    def _fancy_number(x):
        if is_zero(x):
            return 0.0
        else:
            return x


def inverse(m: Matrix):
    if m.row_num() != m.col_num():
        return None
    # 只有方阵有逆矩阵

    n = m.row_num()
    linear_system = LinearSystem(m, Matrix.E(n))
    if not linear_system.gauss_jordan_elimination():
        return None
    else:
        return Matrix([[row[i] for i in range(n, n * 2)]
                       for row in linear_system.Ab])


def rank(m: Matrix):
    ls = LinearSystem(m)
    ls.gauss_jordan_elimination()
    # 生成对应列数的零向量
    zero_vector = Vector.zero(m.col_num())
    # 判断非零行个数
    return sum(zero_vector != a for a in ls.Ab)
